from coin.exchange.base.kr_rest.product import NativeProduct
from coin.exchange.base.kr_rest.futures_product import NativeFuturesProduct
import coin.proto.coin_feed_pb2 as proto_pb2


class LevelElement(object):
  def __init__(self, price, quantity, count):
    self.price = price
    self.quantity = quantity
    self.count = count

  def __str__(self):
    return "<%s,%s,%s>" % (self.price, self.quantity, self.count)

  def as_proto_int(self):
    if self.count is not None and self.count > 0:
      return proto_pb2.BookEntry(price=int(self.price),
                                 qty_float=float(self.quantity),
                                 count=int(self.count))
    else:
      return proto_pb2.BookEntry(price=int(self.price), qty_float=float(self.quantity))

  def as_proto_float(self):
    if self.count is not None and self.count > 0:
      return proto_pb2.BookEntry(price_float=float(self.price),
                                 qty_float=float(self.quantity),
                                 count=int(self.count))
    else:
      return proto_pb2.BookEntry(price_float=float(self.price), qty_float=float(self.quantity))


class LevelBook(object):
  def __init__(self, product):
    assert isinstance(product, (NativeProduct, NativeFuturesProduct))
    self.product = product
    self.bids = []
    self.asks = []
    self.mid_price = None

  def bid0(self):
    return self.bids[0] if len(self.bids) > 0 else None

  def ask0(self):
    return self.asks[0] if len(self.asks) > 0 else None

  def add_bid(self, elem):
    self.bids.append(elem)

  def add_ask(self, elem):
    self.asks.append(elem)

  def finalize(self):
    self.bids = sorted(self.bids, key=lambda elem: -elem.price)
    self.asks = sorted(self.asks, key=lambda elem: elem.price)
    self.mid_price = (self.bids[0].price + self.asks[0].price) / 2

  def as_proto_int(self):
    return proto_pb2.BookMessage(
        ask=[x.as_proto_int() for x in self.asks],
        bid=[x.as_proto_int() for x in self.bids],
    )

  def as_proto_float(self):
    return proto_pb2.BookMessage(
        ask=[x.as_proto_float() for x in self.asks],
        bid=[x.as_proto_float() for x in self.bids],
    )

  def __str__(self):
    return "<LevelBook %s %s %d %d %s %s>" % (
        self.product, self.mid_price, len(self.bids), len(self.asks), self.bid0(), self.ask0())
